import torch
import torchvision.datasets
from torch import nn
from torch.nn import MaxPool2d
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter


dataset = torchvision.datasets.CIFAR10("../../dataSet",
                                       train=False,
                                       download=True,
                                       transform=torchvision.transforms.ToTensor())

dataloader = DataLoader(dataset,64)


#
# input = torch.tensor([[1,2,0,3,1],
#                       [0,1,2,3,1],
#                       [1,2,1,0,0],
#                       [5,2,3,1,1],
#                       [2,1,0,1,1]],dtype=torch.float32)
#
# input = torch.reshape(input,(-1,1,5,5))
#
# print(input.shape)


class Ah(nn.Module):
    def __init__(self):
        super().__init__()
        self.maxpool1 = MaxPool2d(kernel_size=3,ceil_mode=True)

    def forward(self,input):
        output = self.maxpool1(input)
        return output

ah = Ah()
# output = ah(input)
#
# print(output)


step = 0
writer = SummaryWriter("../logs")
for data in dataloader:
    imags ,tables =data
    writer.add_images("pool_input",imags,step)
    output = ah(imags)
    # cihua buhui gaibian chniall
    writer.add_images("pool_output",output,step)
    step +=1